package com.shiyi.shiyiaiagent.chatMemory;

import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.shiyi.shiyiaiagent.chatMemory.entity.ChatMessage;
import com.shiyi.shiyiaiagent.service.ChatMessagesService;
import jakarta.annotation.Resource;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.stereotype.Component;

import java.util.List;

/**
 * 基于MySQL的聊天记忆实现
 */
@Component
public class MySQLChatMemory implements ChatMemory {

  @Resource
  private ChatMessagesService chatMessagesService;
  @Override
  public void add(String conversationId, List<Message> messages) {
    List<ChatMessage> chatMessages = messages.stream().
      map(message -> ChatMessage.fromMessage(conversationId, message))
      .toList();
    chatMessagesService.saveBatch(chatMessages);
  }

  @Override
  public List<Message> get(String conversationId, int lastN) {
    QueryWrapper<ChatMessage> chatMessageQueryWrapper = new QueryWrapper<>();
    chatMessageQueryWrapper.eq("conversation_id", conversationId)
      .orderByAsc("created_at");
    List<ChatMessage> chatMessages = chatMessagesService.list(chatMessageQueryWrapper);
    return chatMessages.stream().skip(Math.max(0, chatMessages.size() - lastN)).map(ChatMessage::toMessage).toList();
  }

  @Override
  public void clear(String conversationId) {
    QueryWrapper<ChatMessage> chatMessageQueryWrapper = new QueryWrapper<>();
    chatMessageQueryWrapper.eq("conversation_id", conversationId);
    chatMessagesService.remove(chatMessageQueryWrapper);
  }


}
